# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

from quantlib import *
from StochasticPathGeneratorMod import *
from SimFileMod import *


def runExample():
    import matplotlib.pyplot as plt
    

    # Initial settings
    todaysDate = Date(15,May,1998)
    saveFileName = 'simPaths'
    nSims = 25

    spg = StochasticPathGenerator()

    # Make a random number generator
    spg.setGaussianRandomNumberGenerator('MersenneTwister-InverseNormal')

    # Choose which processes to simulate

    #   1 - Black-Scholes-Merton
    initial = 100.0
    vol = 0.15
    div = 0.0
    rate = 0.05
    settlementDate = Date(17,May,1998)
    underlying = SimpleQuote(initial)
    volatility = BlackConstantVol(todaysDate, TARGET(), vol, Actual365Fixed())
    dividendYield = FlatForward(settlementDate, div, Actual365Fixed())
    riskFreeRate = FlatForward(settlementDate, rate, Actual365Fixed())
    process1 = BlackScholesMertonProcess(QuoteHandle(underlying),
                                         YieldTermStructureHandle(dividendYield),
                                         YieldTermStructureHandle(riskFreeRate),
                                         BlackVolTermStructureHandle(volatility))
    spg.registerMacroProcess(process1, 'bsm')

    #   2 - Heston
    initial = 100.0
    vol0 = 0.15
    volSigma = 0.01
    volKappa = 0.3
    volTheta = 0.15
    div = 0.0
    rate = 0.05
    rho = 0.0
    settlementDate = Date(17,May,1998)
    underlying = SimpleQuote(initial)
    dividendYield = FlatForward(settlementDate, div, Actual365Fixed())
    riskFreeRate = FlatForward(settlementDate, rate, Actual365Fixed())
    process2 = HestonProcess(YieldTermStructureHandle(riskFreeRate),
                             YieldTermStructureHandle(dividendYield),
                             QuoteHandle(underlying),
                             vol0,
                             volKappa,
                             volTheta,
                             volSigma,
                             rho)
    spg.registerMacroProcess(process2, 'hest')

    # Set some correlations between the processes
    spg.setCorrelation('bsm',0,'hest',0, 0.1)
    spg.setCorrelation('hest',0,'hest',1, -0.3)

    # Choose time points to simulate for
    timeGrid = [0.5, 2.4, 2.89, 4.6]
    spg.initialiseSimulation(timeGrid)

    # Simulate and save to files
    saveFile1Name = saveFileName + '_1.sim'
    saveFile2Name = saveFileName + '_2.sim'
    saveFile1 = SaveSimFile(saveFile1Name, timeGrid)
    saveFile2 = SaveSimFile(saveFile2Name, timeGrid)
    for idx in range(nSims):
        spg.simulateNext()
        saveFile1.writeSim(spg.getPath('bsm',0))
        saveFile2.writeSim(spg.getPath('hest',0))
    saveFile1.close()
    saveFile2.close()

    # Plot paths from saved file
    fig = plt.figure()
    ax = fig.add_subplot(111)
    loadFile1 = LoadSimFile(saveFile1Name)
    xVals = [loadFile1.timePoints(idx) for idx in range(len(timeGrid)+1)]
    for path in loadFile1:
        yVals = path
        ax.plot(xVals, yVals)
    loadFile1.close()
    plt.show()


if __name__ == "__main__":
    runExample()